#!/usr/bin/python
#==========================================================================
# (c) 2007  Total Phase, Inc.
#--------------------------------------------------------------------------
# Project : Cheetah Examples
# File    : flash.py
#--------------------------------------------------------------------------
# Read from, erase, write to, and verify data contents of the High-Speed
# SPI Flash board
#--------------------------------------------------------------------------
# Redistribution and use of this file in source and binary forms, with
# or without modification, are permitted.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE
# COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
#==========================================================================


#==========================================================================
# IMPORTS
#==========================================================================
from cheetah_py import *
from time import *
import os.path

#=========================================================================
# CONSTANTS
#=========================================================================
INTEGRITY_CHECK_NUM_BYTES  =  0x1000000     # When to stop data verify
INTEGRITY_LOOP_SIZE        =  (0x400000/4)  # Flash size divided by 4

# If page size is larger than 1K, _write function will have to be modified
# to take into account the possibility of starting the write in the middle 
# of a page.
PAGE_SIZE                  =  128
PAGE_WRITE_BATCH_SIZE      =  16
PAGE_PROGRAM_CYCLE_TIME_NS =  2100000

COMMAND_READ   = 0
COMMAND_WRITE  = 1
COMMAND_ERASE  = 2
COMMAND_VERIFY = 3
COMMAND_EDI    = 4
COMMAND_ECHIP  = 5
COMMAND_RESET  = 6

#=========================================================================
# FUNCTIONS
#=========================================================================
def _timeMicroseconds():
    return long(time() * 1000000L)

def _printProgress (percent, elapsedTime, blockNum, blocksize): 
    progressbar = " " * 32
    tenths = int(percent/5)
    progressbar = "["
    if(tenths != 0):
        progressbar += tenths * " "
    progressbar += "#"
    if(20-tenths != 0):
        progressbar += (20-tenths) * " "
    progressbar += "]"
    print "\r%s %3d%% %7d KB %.2lf seconds" % \
          (progressbar, percent, blocksize / 1024 * (blockNum+1), elapsedTime),
    sys.stdout.flush()
 
def _connect (port, bitrate, mode):
    # Open the device.
    handle = ch_open(port)
    if (handle <= 0):
        print "Unable to open Cheetah device on port %d" % port
        print "Error code = %d" % handle
        return (1, handle)
    print "Opened Cheetah device on port %d" % port

    ch_host_ifce_speed_string = ""
    
    if (ch_host_ifce_speed(handle)):
        ch_host_ifce_speed_string = "high speed"
    else:
        ch_host_ifce_speed_string = "full speed"
        
    print "Host interface is %s" % ch_host_ifce_speed_string

    # Ensure that the SPI subsystem is configured.
    ch_spi_configure(handle, (mode >> 1), mode & 1, CH_SPI_BITORDER_MSB, 0x0)
    print "SPI configuration set to mode %d, MSB shift, SS[2:0] active low" \
          % mode
    sys.stdout.flush()
    
    # Power the flash using the Cheetah adapter's power supply.
    ch_target_power(handle, CH_TARGET_POWER_ON)
    ch_sleep_ms(100)
    
    # Set the bitrate.
    bitrate = ch_spi_bitrate(handle, bitrate)
    if (bitrate < 0):
        print "Could not set bitrate."
        print "Error code = %d" % bitrate
        return (1, handle)
    
    print "Bitrate set to %d kHz" %  bitrate
    sys.stdout.flush()
    return (0, handle)

def edi_delay(handel,ns_delay):
    ch_spi_queue_clear(handle)
    ch_spi_queue_delay_ns(handle, ns_delay)
    (count, temp)  = ch_spi_batch_shift(handle, 0)
    return (0,0)

def edi_read_byte(handle, addr):

    ch_spi_queue_clear(handle)

    # Enable the cheetah outputs
    ch_spi_queue_oe(handle, 1)
 
    # Set slave select to deasserted state, in case it was left
    # low by a previously interrupted transaction (ctrl-c).  This
    # should reset the state machine.
    ch_spi_queue_ss(handle, 0)
    ch_spi_queue_ss(handle, 1)

    # Queue a EDI read command.
    ch_spi_queue_byte(handle, 1, 0x30)

    # Queue 3 bytes of address.
    ch_spi_queue_byte(handle, 1, (addr >> 16) & 0xff)
    ch_spi_queue_byte(handle, 1, (addr >>  8) & 0xff)
    ch_spi_queue_byte(handle, 1, (addr >>  0) & 0xff)

    # We queue up extra bytes here because the time it takes to 
    # setup another transaction is long enough for the read to complete
    # In that case you don't get a 0x5f or a 0x50
    # I don't want to just delay and then read the data since I think that 
    # will fail at if we start running at faster SPI clock speeds.
    # This way if it fails the fail can be detected.  
    # A simple delay and read has no way of detecting fail.
    ch_spi_queue_byte(handle, 4, 0x00)

    # Send the data
    (count, data_in) = ch_spi_batch_shift(handle, 8)
    
    if count != 8: return (2,count)

    for index, item in enumerate(data_in):
	if item == 0x5f: continue
	if item == 0x50:
	  index+=1
	  break
	      
    ch_spi_queue_ss(handle, 0)

    if index < len(data_in):
      return (0,data_in[index])
    else:
      return (1,0)

def edi_inc_read_byte(handle):

    ch_spi_queue_clear(handle)

    # Enable the cheetah outputs
    ch_spi_queue_oe(handle, 1)
 
    # Set slave select to deasserted state, in case it was left
    # low by a previously interrupted transaction (ctrl-c).  This
    # should reset the state machine.
    ch_spi_queue_ss(handle, 0)
    ch_spi_queue_ss(handle, 1)

    # Queue a EDI read command.
    ch_spi_queue_byte(handle, 1, 0x33)

    # We queue up extra bytes here because the time it takes to 
    # setup another transaction is long enough for the read to complete
    # In that case you don't get a 0x5f or a 0x50
    # I don't want to just delay and then read the data since I think that 
    # will fail at if we start running at faster SPI clock speeds.
    # This way if it fails the fail can be detected.  
    # A simple delay and read has no way of detecting fail.
    ch_spi_queue_byte(handle, 3, 0x00)

    # Send the data
    (count, data_in) = ch_spi_batch_shift(handle, 4)
    
    if count != 4: return (2,count)

    for index, item in enumerate(data_in):
	if item == 0x5f: continue
	if item == 0x50:
	  index+=1
	  break
	      
    ch_spi_queue_ss(handle, 0)

    if index < len(data_in):
      return (0,data_in[index])
    else:
      return (1,0)


def edi_read_rom(handle):

    romdata = array('B')

    start = _timeMicroseconds()
    
    i = 0
    while i < 32*1024:
      (result,data) = edi_read_byte(handle,i)
      i+=1
      if result == 0:
	romdata.append(data)
      else:
	return (-1,romdata)
      # Print out the progress.
      currTime = _timeMicroseconds()
      print "\r%.1f" % ((currTime - start) / 1000000.0),

    return (0,romdata)

def edi_fast_read_rom(handle):

    romdata = array('B')

    start = _timeMicroseconds()
    
    (result,data) = edi_read_byte(handle,0)
    if result == 0:
	romdata.append(data)
    i = 1
    while i < 32*1024:
      (result,data) = edi_inc_read_byte(handle)
      i+=1
      if result == 0:
	romdata.append(data)
      else:
	return (-1,romdata)
      # Print out the progress.
      currTime = _timeMicroseconds()

    return (0,romdata)


def edi_write_byte(handle, addr, byte):
    ch_spi_queue_clear(handle)

    # Enable the cheetah outputs
    ch_spi_queue_oe(handle, 1)
 
    # Set slave select to deasserted state, in case it was left
    # low by a previously interrupted transaction (ctrl-c).  This
    # should reset the state machine.
    ch_spi_queue_ss(handle, 0)
    ch_spi_queue_ss(handle, 1)

    # Queue a EDI write command.
    ch_spi_queue_byte(handle, 1, 0x40)

    # Queue 3 bytes of address.
    ch_spi_queue_byte(handle, 1, (addr >> 16) & 0xff)
    ch_spi_queue_byte(handle, 1, (addr >>  8) & 0xff)
    ch_spi_queue_byte(handle, 1, (addr >>  0) & 0xff)

    # Queue the data to write .
    ch_spi_queue_byte(handle, 1, byte)
    
    # De assert chip select
    ch_spi_queue_ss(handle, 0)
    ch_spi_queue_oe(handle, 0)
    
    # Do it
    (count, data_in) = ch_spi_batch_shift(handle, 0)
    
    if count != 5:
      return (1,count)
      
    return (0,count)

def edi_8051_reset(handle):
    (result,data) = edi_read_byte(handle, 0xf010)
    if result > 0: return (1,0)
    data |= 0x01
    (result,count) = edi_write_byte(handle,0xf010, data)
    if result > 0: return (2,count)
    return (0,0)

def edi_8051_run(handle):
    (result,data) = edi_read_byte(handle, 0xf010)
    if result > 0: return (1,0)
    data &= ~0x01
    (result,count) = edi_write_byte(handle,0xf010, data)
    if result > 0: return (2,count)
    return (0,0)

def edi_flash_select(handle):
    (result,data) = edi_read_byte(handle, 0xf011)
    if result > 0: return (1,0)
    data |= 0x01
    (result,count) = edi_write_byte(handle,0xf011, data)
    if result > 0: return (2,count)
    return (0,0)

def edi_read_id(handle):
    ch_spi_queue_clear(handle)

    # Enable the cheetah outputs
    ch_spi_queue_oe(handle, 1)

    # Set slave select to deasserted state, in case it was left
    # low by a previously interrupted transaction (ctrl-c).  This
    # should reset the state machine.
    ch_spi_queue_ss(handle, 0)
    ch_spi_queue_ss(handle, 1)

    # Queue a EDI read ID.
    ch_spi_queue_byte(handle, 1, 0x3e)

    ch_spi_queue_byte(handle, 1, 0x00)

    # De assert chip select
    ch_spi_queue_ss(handle, 0)
    ch_spi_queue_oe(handle, 0)
    
    # Do it
    (count, data_in) = ch_spi_batch_shift(handle, 2)

    if count != 2:
      return (1,count)
      
    return (0,data_in[1])

def edi_read_trim_byte(handle,addr):
     # A7:A0
    (result,count) = edi_write_byte(handle,0xfea8, (addr & 0xff) )
    if result > 0: return (1,count)

    # A14:A8
    (result,count) = edi_write_byte(handle,0xfea9, ((addr>>8) & 0xff) )
    if result > 0: return (1,count)

    (result,count) = edi_write_byte(handle,0xfea7,0x90)
    if result > 0: return (result,count)

    (result,count) = edi_wait_flash_busy(handle)
    if result > 0: return (result,count)

    (result,data) = edi_read_byte(handle, 0xfeab)
    if result > 0: return (result,count)

    return (result,data)
 
def edi_dump_trim_data(handle):
   print "Trim Data:"
   (result,data) = edi_read_trim_byte(handle,0x100)
   print hex(data),
   (result,data) = edi_read_trim_byte(handle,0x101)
   print hex(data),
   (result,data) = edi_read_trim_byte(handle,0x102)
   print hex(data),
   (result,data) = edi_read_trim_byte(handle,0x1f0)
   print hex(data),
   (result,data) = edi_read_trim_byte(handle,0x1f1)
   print hex(data),
   (result,data) = edi_read_trim_byte(handle,0x1f2)
   print hex(data),
   (result,data) = edi_read_trim_byte(handle,0x1f3)
   print hex(data),
   (result,data) = edi_read_trim_byte(handle,0x1f4)
   print hex(data),
   (result,data) = edi_read_trim_byte(handle,0x1f5)
   print hex(data),
   (result,data) = edi_read_trim_byte(handle,0x1f6)
   print hex(data),
   print

def edi_start(handle):
    # First reset is a dummy to make the EDI port activate
    edi_8051_reset(handle)
    
    (result,count) = edi_8051_reset(handle)
    if result > 0:
	return (result,count)
	
    (result,count) = edi_flash_select(handle)
    if result > 0:
	return (result,count)

    (result,data) = edi_read_id(handle)
    if result > 0:
	return (result,0)
    if data != 3:
	return (1,0)

    edi_dump_trim_data(handle)

    return (0,0)
    
# The "erase" zeros data rather than a normal
# flash erase of 0xff
def edi_erase_page(handle,flash_addr):
    # A7:A0
    (result,count) = edi_write_byte(handle,0xfea8, (flash_addr & 0xff) )
    if result > 0: return (1,count)
    # A14:A8
    (result,count) = edi_write_byte(handle,0xfea9, ((flash_addr>>8) & 0xff) )
    if result > 0: return (1,count)
    # Page erase command 
    (result,count) = edi_write_byte(handle,0xfea7, 0x20 )
    if result > 0: return (1,count)
    return (0,0)

def edi_erase_chip(handle):
    # A7:A0
    (result,count) = edi_write_byte(handle,0xfea8, 0x00 )
    if result > 0: return (1,count)
    # A14:A8
    (result,count) = edi_write_byte(handle,0xfea9, 0x00 )
    if result > 0: return (1,count)
    # Chip erase command 
    (result,count) = edi_write_byte(handle,0xfea7, 0x60 )
    if result > 0: return (1,count)
    return (0,0)
    
def edi_clear_hvpl(handle):
      (result,count) = edi_write_byte(handle,0xfea7,0x80)
      if result > 0: return (1,count)
      return (0,0)
    
def edi_fill_page(handle,addr,buf):
    for byte in buf:
      # A7:A0
      (result,count) = edi_write_byte(handle,0xfea8, (addr & 0xff) )
      if result > 0: return (1,count)
      # A14:A8
      (result,count) = edi_write_byte(handle,0xfea9, ((addr>>8) & 0xff) )
      if result > 0: return (1,count)
      # byte to latch
      (result,count) = edi_write_byte(handle,0xfeaa, byte)
      if result > 0: return (1,count)      
      # Page latch command 
      (result,count) = edi_write_byte(handle,0xfea7, 0x02 )
      if result > 0: return (1,count)
      addr+=1
      
    return (0,0)

def edi_program_page(handle,flash_addr):
    # A7:A0
#    (result,count) = edi_write_byte(handle,0xfea8, (flash_addr & 0xff) )
#    if result > 0: return (1,count)
#    # A14:A8
#    (result,count) = edi_write_byte(handle,0xfea9, ((flash_addr>>8) & 0xff) )
#    if result > 0: return (1,count)
    # Page program command 
    (result,count) = edi_write_byte(handle,0xfea7, 0x70 )
    if result > 0: return (1,count)
    return (0,0)

def edi_wait_flash_busy(handle):
    count = 0
    while True:
      (result,data) = edi_read_byte(handle, 0xfea0)
      if result > 0: return (1,0)
      if (data & 0x80) == 0x80:
	return (0,0)
      count+=1
      if count > 200000:
	return (2,0)

def edi_write_rom(handle, addr, length, data):

    sys.stdout.flush()
    
    start = _timeMicroseconds()

    nr_pages = int( ((length-1) / PAGE_SIZE) +1)

    print "Writing %d pages" % nr_pages

    page = 0
    pstart = addr
    pend   = addr+PAGE_SIZE
    
    while page < nr_pages:
	page_buf = data[pstart:pend]

	(result,count) = edi_wait_flash_busy(handle)
	if result > 0: return (result,count)

	(result,count) = edi_erase_page(handle,pstart)
	if result > 0: return (result,count)

	(result,count) = edi_wait_flash_busy(handle)
	if result > 0: return (result,count)

#	edi_delay(handle, 1000000)
	
	(result,count) = edi_clear_hvpl(handle)
	if result > 0: return (result,count)

	(result,count) = edi_wait_flash_busy(handle)
	if result > 0: return (result,count)

	(result,count) = edi_fill_page(handle,pstart,page_buf)
	if result > 0: return (result,count)

	(result,count) = edi_wait_flash_busy(handle)
	if result > 0: return (result,count)

	(result,count) = edi_program_page(handle,pstart)
	if result > 0: return (result,count)

	(result,count) = edi_wait_flash_busy(handle)
	if result > 0: return (result,count)

	# Print out the progress.
	currTime = _timeMicroseconds()
	pct = int(((float(page)/float(nr_pages))*100))
	if (pct % 20) == 0:
		print "%d%% %.1f" % (pct, (currTime - start) / 1000000.0)

	sys.stdout.flush()

	pstart+=PAGE_SIZE
	pend+=PAGE_SIZE
	page+=1

    return (0,0)

def edi_erase(handle, addr, length):

    sys.stdout.flush()
    
    start = _timeMicroseconds()

    nr_pages = int( ((length-1) / PAGE_SIZE) +1)

    print "Erasing %d pages" % nr_pages

    page = 0
    pstart = addr
    pend   = addr+PAGE_SIZE
    
    while page < nr_pages:
	(result,count) = edi_wait_flash_busy(handle)
	if result > 0: return (result,count)

	(result,count) = edi_erase_page(handle,pstart)
	if result > 0: return (result,count)

#	(result,count) = edi_wait_flash_busy(handle)
#	if result > 0: return (result,count)

#	(result,count) = edi_clear_hvpl(handle)
#	if result > 0: return (result,count)

	# Print out the progress.
	currTime = _timeMicroseconds()
	print "\r%3.0f%% %.1fs" % ((float(page)/float(nr_pages))*100, (currTime - start) / 1000000.0),
	
	sys.stdout.flush()

	pstart+=PAGE_SIZE
	pend+=PAGE_SIZE
	page+=1

    return (0,0)

def verify_images(refdata,testdata):
    i = 0
    integrity_errors = 0
    while i < 32*1024:
       if ( refdata[i] != testdata[i]):
          integrity_errors += 1 
       i+=1
    if integrity_errors > 0:
       return integrity_errors
    else:
       return 0
	
#==========================================================================
# USAGE INFORMATION
#==========================================================================
def print_usage ():
    print \
"""
usage: edi write <filename> 

"""
   
#==========================================================================
# MAIN PROGRAM
#==========================================================================
if (len(sys.argv) < 3 ):
    print_usage()
    sys.exit(1)

port     = 0
bitrate  = 2000
mode     = 0
command  = sys.argv[1]
filename = sys.argv[2]

start_addr = 0
length = 0
    
commandID = -1
error = 0

if(command == "read"):
    commandID = COMMAND_READ

elif (command == "write"):
    commandID = COMMAND_WRITE

elif (command == "erase"):
    commandID = COMMAND_ERASE

elif (command == "verify"):
    commandID = COMMAND_VERIFY
    
elif (command == "edi"):
    commandID = COMMAND_EDI

elif (command == "chiperase"):
    commandID = COMMAND_ECHIP

elif (command == "reset"):
    commandID = COMMAND_RESET
    
else:
    print "Unknown option: %s" % command
    print "Valid options are: read, write, erase, and verify\n"
    print_usage();
    sys.exit(1)

# Connect to and configure the Cheetah.
(ret_code, handle) = _connect(port, bitrate, mode)
if(ret_code):
    sys.exit(1)

if(commandID ==  COMMAND_EDI):
    (result,count) = edi_start(handle)
    if result > 0:
	print "start err", result, count
	
    edi_dump_trim_data(handle)

# Execute the appropriate command.
if(commandID ==  COMMAND_READ):
    (result,count) = edi_start(handle)
    if result > 0:
	print "start err", result, count
	
    (result,romdata) = edi_read_rom(handle)
    if result == 0:
	datafile = open(filename,"wb")
	romdata.tofile(datafile)
    else:
	print "err"
  
elif (commandID == COMMAND_WRITE):
    (result,count) = edi_start(handle)
    if result > 0:
	print "start err", result, count
  	sys.exit(result) 

    filesize = os.path.getsize(filename)
    datafile = open(filename,"rb")
    dataarray = array("B")
    dataarray.fromfile(datafile,filesize)
    (result,count) = edi_write_rom(handle,0,filesize,dataarray)
    if result > 0:
	print "write err", result, count
    print 

    print "Verifying"
    sys.stdout.flush()

    (result,progdata) = edi_fast_read_rom(handle)
    if result != 0:
	print "read err", result
    else:
	result = verify_images(dataarray,progdata)
    if result > 0:
       print
       print "verify **fail** with %d errors" % result
    else:
       print
       print "verifed"
      
elif (commandID == COMMAND_ERASE):
    (result,count) = edi_start(handle)
    if result > 0:
	print "start err", result, count
    
    (result,count) = edi_erase(handle,0x0,0x8000,)
    if result > 0:
	print "Erase error", result, count
	
elif (commandID == COMMAND_VERIFY):
    ret = _verify(handle, int(sys.argv[5]) * 1024, int(sys.argv[6]) * 1024)

elif (commandID == COMMAND_ECHIP):
    (result,count) = edi_start(handle)
    if result > 0:
	print "start err", result, count

    # chip erase command 
    (result,count) = edi_erase_chip(handle)
    if result > 0:
	print "Chip erase error", result, count

elif (commandID == COMMAND_RESET):
    edi_8051_reset(handle)
    edi_8051_reset(handle)

#Turn the CPU back on 
edi_8051_run(handle)

# Close the device.
ch_target_power(handle,CH_TARGET_POWER_OFF)
ch_close(handle);
if (error):
   sys.exit(error)
else:   
   sys.exit(0)
